#!/usr/bin/python -i
#import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import copy

"""
from optparse import OptionParser
parser = OptionParser("usage:  %prog [options] arg1 ")
parser.add_option("-k", "--kmcolumn", dest="kmcolumn")
parser.add_option("-t", "--transient", dest="transient")
"""

# FIXME: Optimize the functions for the calculus of the Force and Potencial.
# XXX: input is delta_r2 = deltaR2 and deltaAxis
def FSFAxis(delta_r2, deltaAxis, sigma2, eps, rc, DrPhiLJOverr_rc):
  if delta_r2 < rc**2:
    return (-DrPhiLJOverr(delta_r2, sigma2, eps) + DrPhiLJOverr_rc )*deltaAxis
  else:
    return 0

def PhiLJ(delta_r2, sigma2, eps):
  ri6  = (sigma2/delta_r2)**3
  ri12 = ri6**2
  return 4*eps*(ri12 - ri6)

def PhiSF(delta_r2, sigma2, eps, rc, PhiLJ_rc, DrPhiLJOverr_rc):
  if delta_r2 < rc**2:
    return PhiLJ(delta_r2, sigma2, eps) - PhiLJ_rc - DrPhiLJOverr_rc*(delta_r2 - rc**2)/2
  else :
    return 0

#DrPhiLJOverr
def DrPhiLJOverr(delta_r2, sigma2, eps):
  # TODO: Make ri6, ri12 global variables.
  ri6  = (sigma2/delta_r2)**3
  ri12 = ri6**2
  return -24*eps*(2*ri12 - ri6)/delta_r2

def FHOAxis(deltaAxis,m,w2):
  return -m*w2*deltaAxis
  
def PhiHO(delta_r2,m,w2):
  return 0.5*m*w2*delta_r2


# MaxwellBoltzman distribution
def FoverGC(x):
  alpha=1.12151
  return alpha*(x**2)*np.exp(-(x-0.5)**2)

def randomMBdist():
  X = 0
  count=0
  while count<1:
    Y=np.random.exponential(1)
    U=np.random.uniform()
    if U <= FoverGC(Y):
      X=Y
      count = count+1
    return X


### Particle class
class Gas:
  def __init__(self,N):
    
    self.N = N
    self.Dim=3
    self.xlen=100
    self.ylen=3 #10
    
    self.halfXlen=0.5*self.xlen
    self.halfYlen=0.5*self.ylen
    
    self.vec_r = np.zeros(self.N*self.Dim).reshape(self.N,self.Dim)
    self.vec_req = np.zeros(self.N*self.Dim).reshape(self.N,self.Dim)
    self.vec_v = np.zeros(self.N*self.Dim).reshape(self.N,self.Dim)
    self.force = np.zeros(self.N*self.Dim).reshape(self.N,self.Dim)
    
    #FIXME flux a energy density is 1 dimensional
    self.fluxDensity = np.zeros(self.N*self.Dim).reshape(self.N,self.Dim) # j_n
    self.enerDensity = np.zeros(self.N) # h_n
    
    self.deltaX  = np.zeros(self.N*self.N).reshape(self.N,self.N)
    self.deltaY  = np.zeros(self.N*self.N).reshape(self.N,self.N)
    self.deltaZ  = np.zeros(self.N*self.N).reshape(self.N,self.N)
    
    self.deltaR2 = np.zeros(self.N*self.N).reshape(self.N,self.N)
    self.Fx = np.zeros(self.N*self.N).reshape(self.N,self.N)
    self.Fy = np.zeros(self.N*self.N).reshape(self.N,self.N)
    self.Fz = np.zeros(self.N*self.N).reshape(self.N,self.N)
    self.halfV = np.zeros(self.N*self.N).reshape(self.N,self.N) # change to halfV_n
    self.Ekin_n = np.zeros(self.N)
    
    
    self.eps   = 0.1
    self.sigma = 1.0
    self.sigma2= self.sigma**2
    self.rc    = 2.5
    self.rc2   = self.rc**2
    
    self.PhiLJ_rc   = PhiLJ(self.rc2, self.sigma2, self.eps)
    self.DrPhiLJOverr_rc = DrPhiLJOverr(self.rc2, self.sigma2, self.eps)
    
    
    delta_r= 1.0
    tempEq = 1.5
    self.iniPos(delta_r)
    self.iniVel(tempEq)
    
    self.l0 = 1.0
    self.Tleft  = tempEq
    self.Tright = tempEq
  def setXYlen(self, xlen, ylen):
    self.xlen = xlen
    self.ylen = ylen
    self.halfXlen = 0.5*self.xlen
    self.halfYlen = 0.5*self.ylen
    
  def iniPos(self, delta_r):
    self.vec_r[0][0] = -self.xlen/2.0 + delta_r
    self.vec_r[0][1] = -self.ylen/2.0 + delta_r
    
    for i in range(1,self.N):
      self.vec_r[i][0] = self.vec_r[i-1][0] + delta_r
      self.vec_r[i][1] = self.vec_r[i-1][1]
      if (self.vec_r[i][0] + delta_r) > self.halfXlen :
        self.vec_r[i][0] = delta_r - self.halfXlen
        self.vec_r[i][1] = self.vec_r[i-1][1] + delta_r
        if self.vec_r[i][1] > self.halfYlen:
          print 'rescale sigma or delta_r', i
          sys.exit()
    self.vec_req = self.vec_r
  
  def iniVel(self, tempEq):
    np.random.seed(10)
    vx = np.sqrt(tempEq)*np.random.normal(0,1,self.N)
    for i in range(self.N):
      self.vec_v[i,0]=vx[i]
      self.vec_v[i,1]=0.0
      self.vec_v[i,2]=0.0
  
    
  def calculateDeltas(self):
    for i in range(self.N):
      for j in range(self.N):
        self.deltaX[i,j] = self.vec_r[i,0] - self.vec_r[j,0]
        self.deltaY[i,j] = self.vec_r[i,1] - self.vec_r[j,1]
        self.deltaZ[i,j] = self.vec_r[i,2] - self.vec_r[j,2]
    self.deltaR2 = self.deltaX**2 + self.deltaY**2 + self.deltaZ**2
    
  def calculateForces(self):
    for i in range(self.N-1):
      for j in range(i+1,self.N):
        # FSFAxis(delta_r2, deltaAxis, sigma, eps, DrPhiLJOverr_rc)
        self.Fx[i][j] = FSFAxis(self.deltaR2[i][j], self.deltaX[i][j], self.sigma, self.eps, self.rc, self.DrPhiLJOverr_rc)
        self.Fx[j][i] = - self.Fx[i][j]
        
        self.Fy[i][j] = FSFAxis(self.deltaR2[i][j], self.deltaY[i][j], self.sigma, self.eps, self.rc, self.DrPhiLJOverr_rc)
        self.Fy[j][i] = - self.Fy[i][j]
        
#        self.Fz[i][j] = 0.0
#        self.Fz[j][i] = 0.0    
    # Total force over each particle
    for i in range(self.N):
      self.force[i][0] = self.Fx[i].sum()
      self.force[i][1] = self.Fy[i].sum()
#      self.force[i][2] = 0.0
  
  def calculateDeltasHO(self):
  # TODO: XXX
    etha = self.vec_r - self.vec_req
    for i in range(self.N-1):
      j=i+1
      self.deltaX[i,j] = self.vec_r[j,0] - self.vec_r[i,0] - self.l0
      #self.deltaY[i,j] = self.vec_r[j,1] - self.vec_r[i,1]
      #self.deltaZ[i,j] = self.vec_r[j,2] - self.vec_r[i,2] 
      
      self.deltaX[j,i] = -self.deltaX[i,j]
      #self.deltaY[j,i] = self.deltaY[i,j]
      #self.deltaZ[j,i] = self.deltaZ[i,j]
    self.deltaR2 = self.deltaX**2 + self.deltaY**2 + self.deltaZ**2
  
  
  def calculateForcesHO(self):
    m  = 1
    w2 = 1
    l0 = 1
    for i in range(self.N-1):
      j = i+1
      self.Fx[i][j] = FHOAxis(self.deltaX[i][j], m, w2)
      self.Fy[i][j] = FHOAxis(self.deltaY[i][j], m, w2)
      self.Fz[i][j] = FHOAxis(self.deltaZ[i][j], m, w2)
      
      self.Fx[j][i] = - self.Fx[i][j]
      self.Fy[j][i] = - self.Fy[i][j]
      self.Fz[j][i] = - self.Fz[i][j]
    for i in range(self.N):
      self.force[i][0] = self.Fx[i].sum()
      self.force[i][1] = self.Fy[i].sum()

  
  def leftHB(self, gamma, Tleft, Wrange,dt):
    # TODO: Find the particles near the left heat bath
    leftIndx  = np.where(self.vec_r[:,0] < -self.halfXlen + Wrange)[0]
    auxrandom = np.random.normal(0,1,len(leftIndx))
    for i in range(len(leftIndx)):
      indx = leftIndx[i]
      self.force[indx][0] += -gamma*self.vec_v[indx][0]
      self.force[indx][0] += np.sqrt(2*dt*Tleft)*auxrandom[i]
        
  def rightHB(self, gamma, Tright, Wrange,dt):
    # TODO: Find the particles near the left heat bath
    rightIndx  = np.where(self.vec_r[:,0] > self.halfXlen - Wrange)[0]
    auxrandom = np.random.normal(0,1,len(rightIndx))
    for i in range(len(rightIndx)):
      indx = rightIndx[i]
      self.force[indx][0] += -gamma*self.vec_v[indx][0]
      self.force[indx][0] += np.sqrt(2*dt*Tright)*auxrandom[i]
        

  def stepPosition(self, old, dt):
    for i in range(self.N):
      self.vec_r[i][0] = old.vec_r[i][0] + dt*old.vec_v[i][0] + 0.5*dt*dt*(old.force[i][0])
      self.vec_r[i][1] = old.vec_r[i][1] + dt*old.vec_v[i][1] + 0.5*dt*dt*(old.force[i][1])
      self.vec_r[i][2] = old.vec_r[i][2] + dt*old.vec_v[i][2] + 0.5*dt*dt*(old.force[i][2])
  
  #--- Validate Particle in the box
  # FIXME: overlap is possible!! How to avoid that?
  def inBox(self, old):
    for i in range(self.N):
      # out of x axis
      if self.vec_r[i][0] < -self.halfXlen:
        delta_x = -self.halfXlen - self.vec_r[i][0]
        #print delta_x,self.vec_v[i][0], self.vec_r[i][0]
        self.vec_r[i][0] = -self.halfXlen + delta_x
        old.vec_r[i][0]  =  self.vec_r[i][0]
        ### LEFT HEAT BATH ---------------------------- TODO
        self.vec_v[i][0] = np.sqrt(self.Tleft)*randomMBdist() #-self.vec_v[i][0]
        old.vec_v[i][0]  = np.sqrt(self.Tleft)*randomMBdist() #-old.vec_v[i][0]
        #print delta_x,self.vec_v[i][0], self.vec_r[i][0], '*'*5,i
      elif self.vec_r[i][0] > self.halfXlen :
        delta_x =  self.vec_r[i][0] - self.halfXlen
        self.vec_r[i][0] =  self.halfXlen - delta_x
        old.vec_r[i][0]  =  self.vec_r[i][0]
        ### RIGHT HEAT BATH ---------------------------- TODO
        self.vec_v[i][0] = -np.sqrt(self.Tright)*randomMBdist() #-self.vec_v[i][0]
        old.vec_v[i][0]  = -np.sqrt(self.Tright)*randomMBdist() #-old.vec_v[i][0]
        
      # out of y axis
      if -self.halfYlen > self.vec_r[i][1]:
        delta_y = -self.halfYlen - self.vec_r[i][1]
        self.vec_r[i][1] = -self.halfYlen + delta_y
        self.vec_v[i][1] = -self.vec_v[i][1]
      elif self.halfYlen < self.vec_r[i][1] :
        delta_y =  self.vec_r[i][1] - self.halfYlen
        self.vec_r[i][1] =  self.halfYlen - delta_y
        self.vec_v[i][1] = -self.vec_v[i][1]
  
  
  def stepVelocity(self, old, dt):
    for i in range(self.N):
      self.vec_v[i][0] = old.vec_v[i][0] + 0.5*dt*(self.force[i][0] + old.force[i][0])
      self.vec_v[i][1] = old.vec_v[i][1] + 0.5*dt*(self.force[i][1] + old.force[i][1])

  
  # Contribution of each particle to the potential energy.
  def calculateHalfV(self):
    for i in range(self.N):
      for j in range(i+1,self.N):
        self.halfV[i][j] = 0.5*( PhiSF(self.deltaR2[i][j], self.sigma2, self.eps, self.rc, self.PhiLJ_rc, self.DrPhiLJOverr_rc))
        self.halfV[j][i] = self.halfV[i][j]
        
  def calculateEkin_n(self):
    for i in range(self.N):
      self.Ekin_n[i] = 0.5*( self.vec_v[i][0]**2 + self.vec_v[i][1]**2 + self.vec_v[i][2]**2 ) 
  
  # h_n  
  def calculateEnerDensity(self):
    for i in range(self.N):
      self.enerDensity[i] = self.halfV[i].sum() + self.Ekin_n[i]
    
  # j_n
  def calculateFluxDensity(self):
    for i in range(self.N):
      aux_vec = np.zeros(self.Dim)
      for j in range(self.N):
        if i != j :
          aux  = self.Fx[i][j]*(self.vec_v[i][0] + self.vec_v[j][0])
          aux += self.Fy[i][j]*(self.vec_v[i][1] + self.vec_v[j][1])
          aux += self.Fz[i][j]*(self.vec_v[i][2] + self.vec_v[j][2])
          aux_vec = ( self.vec_r[i] - self.vec_r[j] ) * aux
      self.fluxDensity[i] = self.enerDensity[i]*self.vec_v[i] + 0.5*(aux_vec)
  
  def plot(self):
    my_dpi=96
    plt.rc('font', size=25)
    plt.rc('text', usetex=True)
    fig = plt.figure(figsize=(800/my_dpi, 800/my_dpi), dpi=my_dpi)
    ax = fig.add_subplot(1, 1, 1)
    ax.set_xlim((-1.0 - self.halfXlen, 1.0 + self.halfXlen))
    ax.set_ylim((-1.0 - self.halfYlen, 1.0 + self.halfYlen))
    for i in range(self.N):
      circle=plt.Circle((self.vec_r[i][0], self.vec_r[i][1]),0.5,color='r',alpha=0.5)
      ax.add_patch(circle)
    
    ax.set_aspect('equal')
    ax.quiver(self.vec_r[:,0], self.vec_r[:,1], self.vec_v[:,0], self.vec_v[:,1], color='b', pivot='mid', alpha =0.5, units='inches')
    plt.scatter(self.vec_r[:,0], self.vec_r[:,1])
    plt.xlabel("$x$")
    plt.ylabel("$y$")
    plt.grid()
    
#    plt.savefig("figura.pdf",dpi=100)
    plt.show()
    
  def savetxt(self,filename):
    np.savetxt(filename, np.hstack((self.vec_r,self.vec_v)) )
    



##################################################################
### GasTime class  
class GasTime:
  def __init__(self, gas, timeDim=1):
    self.timeDim = timeDim
    self.time  = np.zeros(self.timeDim)
    self.Ekin_nt = np.zeros([timeDim, gas.N])
    self.enerDenst = np.zeros([timeDim, gas.N]) #h_nt
    self.fluxDenst = np.zeros([timeDim, gas.N]) #j_nt
    # could be pressure
    self.time_i = 0
    
    #Micro states in time
    self.xt  = np.zeros([timeDim, gas.N])
    self.yt  = np.zeros([timeDim, gas.N])
    self.vxt = np.zeros([timeDim, gas.N])
    self.vyt = np.zeros([timeDim, gas.N])
    
  def saveState(self, gas, step, dt):
    self.time[self.time_i]  = dt*step
    self.Ekin_nt[self.time_i] = gas.Ekin_n
    self.enerDenst[self.time_i] = gas.enerDensity #h_nt
    self.fluxDenst[self.time_i] = gas.fluxDensity[:,0] #j_nt
    self.time_i = self.time_i + 1
    sys.stdout.write(" "+str(self.time_i))
  
  # Call first the function bellow
  def saveMicroState(self, gas, step, dt):
    for j in range(gas.N):
      self.xt[self.time_i][j]  = gas.vec_r[j][0]
      self.yt[self.time_i][j]  = gas.vec_r[j][1]
      self.vxt[self.time_i][j] = gas.vec_v[j][0]
      self.vyt[self.time_i][j] = gas.vec_v[j][0]

  def savetxt(self,filename):
    #TODO: save in a file the gasTime object
    print filename
    np.savetxt(filename, np.hstack((self.time.reshape(-1,1), self.enerDenst,self.fluxDenst)))
    #f_handle = file(filename,'w')
    #np.savetxt(f_handle, data)
    #f_handle.close()
      
#############################################3
"""
gas_old = Gas(5)
gas_new = Gas(5)

gas_ini = copy.deepcopy(gas_old)

gamma  = 1.0
Tleft  = 1.0
Tright = 2.0
Wrange = 1.0
dt = 0.001


######
gas_old.calculateDeltas()
gas_old.calculateForces()
gas_old.leftHB(gamma, Tleft, Wrange)
gas_old.rightHB(gamma, Tright, Wrange)

gas_new.stepPosition(gas_old,dt)
gas_new.inBox()

gas_new.calculateDeltas()
gas_new.calculateForces()
gas_new.stepVelocity(gas_old, dt)

gas_new.calculateHalfV()
gas_new.calculateEnerDensity()  # h_n
gas_new.calculateFluxDensity()  # j_n

#gas_new.calculateEpot()
#gas_new.calculateEkin()
#gas_new.calculateFlux()

#save the values for the next step
gas_old = copy.deepcopy(gas_new)

######


gas_end = copy.deepcopy(gas_new)
"""
